import os
import numpy as np
import torch
from PIL import Image


#####################################
# 类：从输入的实例掩码中提取边界框
"""
代码功能说明：
AppleDataset 类：
初始化：
root_dir：数据集的根目录，包含 images 和 masks 子目录。
transforms：数据转换流程，用于对图像和目标进行预处理。
加载并排序图像和掩码文件，确保它们按顺序排列。
__getitem__ 方法：
加载指定索引的图像和掩码。
将掩码图像转换为 NumPy 数组，并提取唯一对象 ID。
移除背景 ID，将颜色编码的掩码拆分为二进制掩码。
遍历每个掩码，计算边界框坐标，并确保边界框的宽和高不为 0。
将所有内容转换为 PyTorch 张量。
构建目标字典，包含边界框、标签、掩码、图像 ID、面积和人群标志。
应用数据转换（如果有）。
__len__ 方法：
返回数据集的大小。
get_img_name 方法：
获取指定索引的图像名称。
希望这些注释能够帮助你更好地理解代码的功能和实现细节。
"""
#####################################
class AppleDataset(object):
    def __init__(self, root_dir, transforms):
        # 初始化数据集
        self.root_dir = root_dir  # 数据集根目录
        self.transforms = transforms  # 数据转换流程

        # 加载所有图像和掩码文件，确保它们按顺序排列
        self.imgs = list(sorted(os.listdir(os.path.join(root_dir, "images"))))
        self.masks = list(sorted(os.listdir(os.path.join(root_dir, "masks"))))

    def __getitem__(self, idx):
        # 加载图像和掩码
        img_path = os.path.join(self.root_dir, "images", self.imgs[idx])
        mask_path = os.path.join(self.root_dir, "masks", self.masks[idx])

        # 打开图像并转换为 RGB 格式
        img = Image.open(img_path).convert("RGB")
        # 打开掩码图像（每个颜色对应一个实例，0 表示背景）
        mask = Image.open(mask_path)

        # 将 PIL 图像转换为 NumPy 数组
        mask = np.array(mask)
        # 获取掩码中的唯一对象 ID
        obj_ids = np.unique(mask)

        # 移除背景 ID
        obj_ids = obj_ids[1:]

        # 将颜色编码的掩码拆分为一组二进制掩码
        masks = mask == obj_ids[:, None, None]

        # 获取每个掩码的边界框坐标
        num_objs = len(obj_ids)
        boxes = []
        h, w = mask.shape
        for ii in range(num_objs):
            # 获取掩码的非零位置
            pos = np.where(masks[ii])
            # 计算边界框的坐标
            xmin = np.min(pos[1])
            xmax = np.max(pos[1])
            ymin = np.min(pos[0])
            ymax = np.max(pos[0])

            # 如果边界框的宽或高为 0，则跳过
            if xmin == xmax or ymin == ymax:
                continue

            # 限制边界框的坐标范围
            xmin = np.clip(xmin, a_min=0, a_max=w)
            xmax = np.clip(xmax, a_min=0, a_max=w)
            ymin = np.clip(ymin, a_min=0, a_max=h)
            ymax = np.clip(ymax, a_min=0, a_max=h)
            boxes.append([xmin, ymin, xmax, ymax])

        # 将所有内容转换为 PyTorch 张量
        boxes = torch.as_tensor(boxes, dtype=torch.float32)

        # 只有一个类别（苹果）
        labels = torch.ones((num_objs,), dtype=torch.int64)
        masks = torch.as_tensor(masks, dtype=torch.uint8)

        # 图像 ID
        image_id = torch.tensor([idx])
        # 计算每个边界框的面积
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

        # 所有实例都不是人群
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        # 构建目标字典
        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

        # 应用数据转换
        if self.transforms is not None:
            img, target = self.transforms(img, target)

        return img, target

    def __len__(self):
        # 返回数据集的大小
        return len(self.imgs)

    def get_img_name(self, idx):
        # 获取图像名称
        return self.imgs[idx]